/*
 * Copyright (c) 2012, 2013, Oracle and/or its affiliates. All rights reserved.
 * DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.
 *
 * This code is free software; you can redistribute it and/or modify it
 * under the terms of the GNU General Public License version 2 only, as
 * published by the Free Software Foundation.  Oracle designates this
 * particular file as subject to the "Classpath" exception as provided
 * by Oracle in the LICENSE file that accompanied this code.
 *
 * This code is distributed in the hope that it will be useful, but WITHOUT
 * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
 * FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License
 * version 2 for more details (a copy is included in the LICENSE file that
 * accompanied this code).
 *
 * You should have received a copy of the GNU General Public License version
 * 2 along with this work; if not, write to the Free Software Foundation,
 * Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301 USA.
 *
 * Please contact Oracle, 500 Oracle Parkway, Redwood Shores, CA 94065 USA
 * or visit www.oracle.com if you need additional information or have any
 * questions.
 */
package com.i360day.invoker.socket;

import com.i360day.invoker.DefaultMessageReply;
import com.i360day.invoker.HttpInvokerClientFactoryBean;
import com.i360day.invoker.annotation.Argument;
import com.i360day.invoker.common.RemoteModuleUtils;
import com.i360day.invoker.properties.HttpInvokerProperties;
import com.i360day.invoker.socket.handler.ReplyMessageListener;
import com.i360day.invoker.socket.handler.SimpleMessageByteBufferHandler;
import jakarta.websocket.*;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.BeansException;
import org.springframework.context.ApplicationContext;
import org.springframework.context.ApplicationContextAware;
import org.springframework.util.Assert;

import java.math.BigDecimal;
import java.net.URI;
import java.util.*;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.Executor;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.locks.Lock;
import java.util.concurrent.locks.ReentrantLock;
import java.util.function.Consumer;

/**
 * @author liju.z
 * @date 2024/4/3 23:17
 */
public class SimpleWebSocketTemplate implements ApplicationContextAware {
    private Logger logger = LoggerFactory.getLogger(SimpleWebSocketTemplate.class);
    private Lock lock = new ReentrantLock();
    private ApplicationContext applicationContext;
    private WebSocketContainer webSocketContainer;
    private HttpInvokerProperties httpInvokerProperties;
    private Executor executor = command -> new Thread(command).start();
    private Map<String, List<Session>> sessionMap = new ConcurrentHashMap<>();
    private Map<String, ReplyMessageListener> messageHandlerListenerMap = new ConcurrentHashMap<>();

    public SimpleWebSocketTemplate(HttpInvokerProperties httpInvokerProperties) {
        this.httpInvokerProperties = httpInvokerProperties;
        this.webSocketContainer = ContainerProvider.getWebSocketContainer();
    }

    public SimpleWebSocketTemplate(Executor executor, HttpInvokerProperties httpInvokerProperties) {
        this.executor = executor;
        this.webSocketContainer = ContainerProvider.getWebSocketContainer();
        this.httpInvokerProperties = httpInvokerProperties;

    }

    /**
     * 执行
     *
     * @param address
     * @param consumer
     */
    public DefaultMessageReply replySendMessageToServer(URI address, String replyId, Consumer<Session> consumer) {
        DefaultMessageReply messageReply = new DefaultMessageReply(replyId);

        //获取当前链接地址下的所有session
        List<Session> sessionList = Optional.ofNullable(sessionMap.get(address.getAuthority())).orElse(Collections.emptyList());
        Assert.isTrue(!sessionList.isEmpty(), String.format("%s session is null, socket Unlinked", address));

        //随机获取一个session
        Session session = sessionList.get(BigDecimal.valueOf(Math.random() * sessionList.size()).intValue());
        consumer.accept(session);

        //设置回调
        messageHandlerListenerMap.put(messageReply.getReplyTo(), messageReply);
        return messageReply;
    }

    /**
     * 执行消息处理监听
     *
     * @param replyId
     * @param consumer
     */
    public void executeMessageHandlerListener(String replyId, Consumer<ReplyMessageListener> consumer) {
        //获取回调监听
        ReplyMessageListener replyMessageListener = messageHandlerListenerMap.get(replyId);

        //执行通知回调
        Optional.ofNullable(replyMessageListener).ifPresent(listener -> {
            try {
                consumer.accept(listener);
            } finally {
                messageHandlerListenerMap.remove(replyId);
            }
        });
    }

    /**
     * 添加socket
     *
     * @param session
     */
    public void addSocket(Session session) {
        List<Session> sessionList = Optional.ofNullable(sessionMap.get(session.getRequestURI().getAuthority())).orElse(new Vector<>());
        sessionList.add(session);
        sessionMap.put(session.getRequestURI().getAuthority(), sessionList);
    }

    /**
     * 根据地址获取session列表
     *
     * @param address
     * @return
     */
    public List<Session> getAddressSocket(URI address) {
        return sessionMap.get(address.getAuthority());
    }

    /**
     * 删除socket
     *
     * @param session
     */
    public void removeSocket(Session session) {
        List<Session> clientSessionList = Optional.ofNullable(sessionMap.get(session.getRequestURI().getAuthority())).orElse(Collections.emptyList());
        Iterator<Session> iterator = clientSessionList.iterator();
        while (iterator.hasNext()) {
            Session s = iterator.next();
            if (s.equals(session)) {
                if (s.isOpen()) {
                    try {
                        s.close();
                    } catch (Exception ex) {
                        //ignore
                    }
                }
                iterator.remove();
            }
        }
    }

    public ApplicationContext getApplicationContext() {
        return applicationContext;
    }

    @Override
    public void setApplicationContext(ApplicationContext applicationContext) throws BeansException {
        this.applicationContext = applicationContext;

        Map<String, HttpInvokerClientFactoryBean> beansOfType = applicationContext.getBeansOfType(HttpInvokerClientFactoryBean.class);
        if (!beansOfType.values().isEmpty()) {
            for (HttpInvokerClientFactoryBean clientFactoryBean : beansOfType.values()) {
                Class<?> serviceInterface = clientFactoryBean.getServiceInterface();
                Argument clientArgument = RemoteModuleUtils.getRemoteClientArgument(serviceInterface, "createSocketServer");
                if (clientArgument == null || clientArgument.value().equals("true")) {
                    createSocketClient(URI.create(clientFactoryBean.getServiceUrl()));
                }
            }
        }
    }

    /**
     * 客户端链接socket，当socket关闭后自动重连
     *
     * @param address
     */
    private void createSocketClient(URI address) {
        //reset
        executor.execute(() -> {
            while (true && !Thread.interrupted()) {
                try {
                    lock.lock();
                    //已有一个链接存在，无需重复添加
                    List<Session> sessionList = getAddressSocket(address);
                    if (sessionList != null && !sessionList.isEmpty()) {
                        return;
                    }

                    Session session = createSocketConnection(address).get(httpInvokerProperties.getConnectTimeout(), TimeUnit.MILLISECONDS);
                    addSocket(session);
                    return;
                } catch (Exception e) {
                    try {
                        Thread.sleep(1000 * 3);
                    } catch (InterruptedException ex) {
                        throw new RuntimeException(ex);
                    }
                } finally {
                    lock.unlock();
                }
            }
        });
    }

    /**
     * 创建socket链接
     *
     * @param address
     * @return
     */
    private CompletableFuture<Session> createSocketConnection(URI address) {
        CompletableFuture<Session> future = new CompletableFuture<>();
        try {
            Endpoint endpoint = new Endpoint() {
                @Override
                public void onOpen(Session session, EndpointConfig config) {
                    System.out.println("onOpen");
                    future.complete(session);
                    executor.execute(() -> {
                        try {
                            while (session.isOpen() && !Thread.interrupted()) {
                                session.getAsyncRemote().sendText("ping");
                                Thread.sleep(1000 * 3);
                            }
                        } catch (Exception ex) {
                            throw new IllegalThreadStateException(ex.getMessage());
                        }
                    });
                }

                @Override
                public void onClose(Session session, CloseReason closeReason) {
                    logger.info(" {} client {} onClose", session.getRequestURI(), session.getUserProperties().get("Client-Id"));
                    try {
                        //remote
                        removeSocket(session);

                        //close
                        if (session.isOpen()) {
                            session.close(closeReason);
                        }

                        //reset
                        createSocketClient(address);
                    } catch (Exception ex) {
                        logger.error(" {} client {} onClose error {}", session.getRequestURI(), session.getUserProperties().get("Client-Id"), ex);
                    }
                }

                @Override
                public void onError(Session session, Throwable thr) {
                    logger.error(" {} client {} onError {}", session.getRequestURI(), session.getUserProperties().get("Client-Id"), thr);
                }
            };

            //configure
            String clientId = UUID.randomUUID().toString();
            ClientEndpointConfig.Builder builder = ClientEndpointConfig.Builder.create();
            builder.configurator(new ClientEndpointConfig.Configurator() {
                @Override
                public void beforeRequest(Map<String, List<String>> headers) {
                    headers.put("Client-Id", Arrays.asList(clientId));
                }
            });
            ClientEndpointConfig clientEndpointConfig = builder.build();
            Map<String, Object> userProperties = clientEndpointConfig.getUserProperties();
            userProperties.put("Client-Id", clientId);

            //create Connection
            String wsUrl = String.format("ws://%s%s", address.getAuthority(), httpInvokerProperties.getWebSocket().getEndpoint());
            Session session = webSocketContainer.connectToServer(endpoint, clientEndpointConfig, URI.create(wsUrl));

            session.addMessageHandler(new SimpleMessageByteBufferHandler(endpoint, session, this));
        } catch (Exception ex) {
            future.completeExceptionally(ex);
        }
        return future;
    }
}
